"""
@FileName：dataset.py
@Description：定义数据集
@Author：Lemon
@Time：2025/5/8 15:44
"""
import os

from PIL import Image
import torch
import torchvision

from process import get_dirs, get_files


def load(img_path):
    types = ['Northwind', 'FreeForm']
    paths = []
    for t in types:
        dirs = get_dirs(os.path.join(img_path, t))
        for d in dirs:
            files = get_files(d)
            for file in files:
                paths.append(os.path.join(d, file))
    return paths


class MyDataset(torch.utils.data.Dataset):
    def __init__(self, img_paths, label_map):
        self.img_paths = load(img_paths)
        self.label_map = label_map
        self.transform = torchvision.transforms.Compose([
            torchvision.transforms.Resize((128, 128)),
            torchvision.transforms.ToTensor(),
            torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, item):
        path = self.img_paths[item]
        img = Image.open(self.img_paths[item])
        img = self.transform(img)
        file_id = os.path.basename(os.path.dirname(path))  # 修改这里
        label = self.label_map[file_id]
        return img, label, path
